// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/vector/base_vector.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux {
    namespace exec {
        class EvalCtx;
        class EvalErrors;

        using EvalErrorsPtr = std::shared_ptr<EvalErrors>;
    } // namespace exec

    // Represents a function with possible captures.
    class Callable {
    public:
        virtual ~Callable() = default;

        virtual bool hasCapture() const = 0;

        /// Applies 'this' to 'args' for 'rows' and returns the result in
        /// '*result'.
        /// @param rows The rows that this callable applies to. It is the element rows
        /// of the complex-typed input of the lambda function.
        /// @param wrapCapture A mapping that translates row numbers in 'rows' to the
        /// corresponding numbers for captured variables, i.e. is an indices vector
        /// that is used for wrapping captured variables in a dictionary before
        /// passing these to the function. The typical use case of lambdas applies a
        /// function to elements of repeated types, so that the values of the
        /// arguments and captures are not aligned. This serves to align these. If
        /// nullptr, the captures are passed as is.
        /// @param validRowsInReusedResult This selectivity vector is used to store
        /// all valid rows in 'result' that can be reused between multiple Callables
        /// in a function vector. It helps preserve rows that are valid in 'result'
        /// but on which this callable does not apply.
        /// @param elementToTopLevelRows A mapping from element rows (i.e., the rows
        /// argument) to top-level rows of the complex-typed input of the lambda
        /// function. elementToTopLevelRows could be a nullptr, meaning the element
        /// rows are identical to top-level rows.
        virtual void apply(
            const SelectivityVector &rows,
            const SelectivityVector *validRowsInReusedResult,
            const BufferPtr &wrapCapture,
            exec::EvalCtx *context,
            const std::vector<VectorPtr> &args,
            const BufferPtr &elementToTopLevelRows,
            VectorPtr *result) = 0;

        /// Same as 'apply', but errors are suppressed and returned in
        /// 'elementErrors', and errors in 'context' are not updated.
        virtual void applyNoThrow(
            const SelectivityVector &rows,
            const SelectivityVector *validRowsInReusedResult,
            const BufferPtr &wrapCapture,
            exec::EvalCtx *context,
            const std::vector<VectorPtr> &args,
            exec::EvalErrorsPtr &elementErrors,
            VectorPtr *result) = 0;
    };

    // Represents a vector of functions. In most cases all the positions
    // in the vector have the same function. We allow for different
    // functions at different positions for completeness.
    class FunctionVector : public BaseVector {
    public:
        // Iterates over distinct functions that fall inside 'rows'. The
        // Callable and the applicable subset of 'rows' is returned for each
        // distinct function by next(). A vectorized lambda application
        // first loops over the functions, then applies each to the
        // applicable rows. In most situations, there is only one function in
        // the FunctionVector.
        class Iterator {
        public:
            struct Entry {
                /// Callable lambda.
                Callable *callable;

                /// Rows that lambda applies to.
                SelectivityVector *rows;

                operator bool() const {
                    return callable != nullptr;
                }
            };

            Iterator(const FunctionVector *vector, const SelectivityVector *rows)
                : rows_(*rows),
                  functions_{vector->functions_},
                  rowSets_{vector->rowSets_} {
            }

            Entry next() {
                while (index_ < functions_.size()) {
                    effectiveRows_ = rowSets_[index_];
                    effectiveRows_.intersect(rows_);
                    if (!effectiveRows_.hasSelections()) {
                        ++index_;
                        continue;
                    }
                    Entry entry{functions_[index_].get(), &effectiveRows_};
                    ++index_;
                    return entry;
                }
                return {nullptr, nullptr};
            }

        private:
            const SelectivityVector &rows_;
            const std::vector<std::shared_ptr<Callable> > &functions_;
            const std::vector<SelectivityVector> &rowSets_;
            int32_t index_ = 0;
            SelectivityVector effectiveRows_;
        };

        FunctionVector(pollux::memory::MemoryPool *pool, TypePtr type)
            : BaseVector(
                pool,
                type,
                VectorEncoding::Simple::FUNCTION,
                BufferPtr(nullptr),
                0) {
        }

        // Implements evaluation of the lambda function literal special
        // form. This assigns a function to a specified set of rows. This
        // supports a situation of returning different lambdas from
        // different cases of a conditional. FunctionVectors enforce single
        // assignment. A lambda cannot be for example a loop variable.
        void addFunction(
            std::shared_ptr<Callable> callable,
            const SelectivityVector &rows) {
            for (auto &otherRows: rowSets_) {
                auto begin = std::max(rows.begin(), otherRows.begin());
                auto end = std::min(rows.end(), otherRows.end());
                POLLUX_CHECK(
                    !bits::hasIntersection(
                        otherRows.asRange().bits(), rows.asRange().bits(), begin, end),
                    "Functions in a FunctionVector may not have intersecting SelectivityVectors");
            }

            rowSets_.push_back(rows);
            functions_.push_back(callable);
        }

        bool contains_null_at(vector_size_t idx) const override {
            return false;
        }

        std::optional<int32_t> compare(
            const BaseVector * /*other*/,
            vector_size_t /*index*/,
            vector_size_t /*otherIndex*/,
            CompareFlags /*flags*/) const override {
            throw std::logic_error("compare not defined for FunctionVector");
        }

        uint64_t hash_value_at(vector_size_t /*index*/) const override {
            throw std::logic_error("not defined for FunctionVector");
        }

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override {
            throw std::logic_error("hash_all not defined for FunctionVector");
        }

        const BaseVector *loaded_vector() const override {
            throw std::logic_error("loaded_vector not defined for FunctionVector");
        }

        const BaseVector *wrapped_vector() const override {
            return this;
        }

        vector_size_t wrapped_index(vector_size_t index) const override {
            return index;
        }

        Iterator iterator(const SelectivityVector *rows) const {
            return Iterator(this, rows);
        }

        VectorPtr slice(vector_size_t, vector_size_t) const override {
            POLLUX_NYI();
        }

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool * /* pool */ = nullptr) const override {
            POLLUX_UNSUPPORTED("copy_preserve_encodings not defined for FunctionVector");
        }

    private:
        std::vector<std::shared_ptr<Callable> > functions_;
        std::vector<SelectivityVector> rowSets_;
    };
} // namespace kumo::pollux
